//! 这个模块是参考golang的singleflight实现的多线程消重实现.
use std::hash::Hash;
use std::{fmt::Debug, sync::Mutex, collections::HashMap};
use std::sync::Arc;

use crossbeam::sync::ShardedLock;

/// 一次消重调用的控制对象.
struct Call<T>
{
    /// call执行的结果, 注意由于Arc由于指针有效性的问题,
    /// 所有有效的Arc都是非零的, 
    /// 所以[`Option<Arc<T>>`] 和 [`Arc<T>`]的空间大小是一样的.
    /// ```rust
    /// use std::mem::size_of;
    /// use std::sync::Arc;
    /// assert_eq!(size_of::<Arc<String>>(), size_of::<Option<Arc<String>>>());
    /// ```
    res: ShardedLock<Option<Arc<T>>>,
}

impl<T> Call<T>
{
    /// 创建一个空的控制对象, 结果初始化为空.
    fn new() -> Self {
        Self { 
            res: Default::default()
        }
    }
}

/// 用于多重调用消除的控制组.
pub struct Group<K, T> 
    where K: Hash + Eq + Clone,
{
    /// 调用的控制表
    m: Mutex<HashMap<K, Arc<Call<T>>>>
}

impl<K, T> Default for Group<K, T> 
    where T: Debug,
          K: Hash + Eq + Clone,
{
    fn default() -> Self {
        Self { m: Default::default() }
    }
}

impl<K, T> Group<K, T> 
    where T: Debug,
          K: Hash + Eq + Clone,
{

    /// 创建一个新的控制组.
    pub fn new() -> Self {
        Self { 
            m: Mutex::new(HashMap::new())
        }
    }

    pub fn call_arc<F>(&self, key: &K, func: F) -> Arc<T>
    where
        F: FnOnce() -> Arc<T>
    {
        // 第一阶段, 锁表, 获取call对象, 这一阶段的优化点在于尽可能减少锁住控制表的时间.
        let mut m = self.m.lock().unwrap();
        if let Some(c) = m.get(key) {
            // 控制表内部存在call对象, 复制call对象(浅拷贝).
            let c = c.clone();
            drop(m);
            // 第一阶段结束, 释放控制表锁.

            // 第二阶段, 读写锁取读锁. 等待call对象完成
            let out = c.res.read().unwrap();

            // 注意这里的unwrap. 由于唯一一次赋值赋值的是Some.
            // 所以unwrap一定不会报错
            return out.clone().unwrap()
        }

        // 如果控制表内没有call对象, 生成call对象
        let c = Arc::new(Call::new());
        let entry = c.clone();
        // 预先锁定call对象中的读写锁的写锁. 使用读写锁作为线程间的控制机制.
        let mut working_lock = c.res.write().unwrap();
        m.insert(key.clone(), entry); // 将创建的call对象放进控制表中.
        drop(m);
        // 第一阶段结束, 释放控制表锁.

        // 第二阶段, 执行函数, 获取结果.
        let ret = func();
        let res = Some(ret.clone());
        *working_lock = res; // 更新call对象的内容. 注意res一定是Some
        drop(working_lock); // 执行完成, 释放写锁, 允许其他线程读取结果.

        // 收尾, 重新获取控制表的锁. 删除call对象.
        // 由于不同key之间不会互相影响,所以只考虑单一key的情况.
        // 只有控制表中没有call对象的情况下, 才能创建新的call对象,
        // 所以在当前线程删除call对象之前, 不会存在非当前线程的创建的call对象.
        // 所以一定不会删除非当前线程创建的call对象.
        let mut m = self.m.lock().unwrap();
        m.remove(key).unwrap();
        drop(m);
        ret
    }

    /// 进行一次调用, 同一时间, 同一个`key`, 只会有一个实际调用运行.
    pub fn call<F>(&self, key: &K, func: F) -> Arc<T>
    where
        F: FnOnce() -> T,
    {
        // 第一阶段, 锁表, 获取call对象, 这一阶段的优化点在于尽可能减少锁住控制表的时间.
        let mut m = self.m.lock().unwrap();
        if let Some(c) = m.get(key) {
            // 控制表内部存在call对象, 复制call对象(浅拷贝).
            let c = c.clone();
            drop(m);
            // 第一阶段结束, 释放控制表锁.

            // 第二阶段, 读写锁取读锁. 等待call对象完成
            let out = c.res.read().unwrap();

            // 注意这里的unwrap. 由于唯一一次赋值赋值的是Some.
            // 所以unwrap一定不会报错
            return out.clone().unwrap()
        }

        // 如果控制表内没有call对象, 生成call对象
        let c = Arc::new(Call::new());
        let entry = c.clone();
        // 预先锁定call对象中的读写锁的写锁. 使用读写锁作为线程间的控制机制.
        let mut working_lock = c.res.write().unwrap();
        m.insert(key.clone(), entry); // 将创建的call对象放进控制表中.
        drop(m);
        // 第一阶段结束, 释放控制表锁.

        // 第二阶段, 执行函数, 获取结果.
        let ret = Arc::new(func());
        let res = Some(ret.clone());
        *working_lock = res; // 更新call对象的内容. 注意res一定是Some
        drop(working_lock); // 执行完成, 释放写锁, 允许其他线程读取结果.

        // 收尾, 重新获取控制表的锁. 删除call对象.
        // 由于不同key之间不会互相影响,所以只考虑单一key的情况.
        // 只有控制表中没有call对象的情况下, 才能创建新的call对象,
        // 所以在当前线程删除call对象之前, 不会存在非当前线程的创建的call对象.
        // 所以一定不会删除非当前线程创建的call对象.
        let mut m = self.m.lock().unwrap();
        m.remove(key).unwrap();
        drop(m);
        ret
    }

    /// 返回group是否是空的
    pub fn is_empty(&self) -> bool {
        self.m.lock().unwrap().is_empty()
    }
}

#[cfg(test)]
mod test {
    use std::sync::atomic::Ordering;
    use super::*;

    /// 使用多个线程测试[`Group`]的正确性.
    #[test]
    fn test_multiple_threads() {
        use std::time::Duration;
        use crossbeam::scope;

        // 昂贵函数的休眠时间.
        const SLEEP_DURATION: Duration = Duration::from_secs(5);
        // 昂贵函数的返回值.
        const RES: usize = 10;

        // 测试用的线程数.
        const THREAD_COUNT: u32 = 40;
        // 测试使用的key数量.
        const KEY_COUNT: u32 = 8;
        // 测试的线程数一定要大于key数量, 否则key数就失去了意义.
        assert!(THREAD_COUNT >= KEY_COUNT);

        // 昂贵函数执行次数的计数器.
        let count: std::sync::atomic::AtomicU32 = 0.into();

        // 一个模拟的昂贵函数.
        let expensive_fn = || {
            println!("begin expensive fn");
            std::thread::sleep(SLEEP_DURATION);
            println!("finish expensive fn");
            count.fetch_add(1, Ordering::Relaxed);
            RES
        };

        let control_group: Group<u32, usize> = Group::new();

        let start = std::time::Instant::now();
        // 使用scoped thread进行测试.
        scope(|s| {
            // 启动一定数量的thread.
            for index in 0 .. THREAD_COUNT {
                // 计算出thread所属的key.
                let key = index % KEY_COUNT;
                // 借用control_group, 这样只有借用被移动到闭包中.
                let g = &control_group;
                s.spawn(move |_| {
                    // 用key为键进行调用, 因为线程启动速度非常快, 所以同一key的多个线程应该只会实际调用一次昂贵函数.
                    let res = g.call(&key, expensive_fn);
                    // 返回结果应当是正确的.
                    assert_eq!(*res, RES);
                });
            }
        }).unwrap();
        let elapsed = start.elapsed();

        // 由于所有的线程是一起启动的, 所以expensive_fn应该只被调用了KEY_COUNT次.
        assert!(count.load(Ordering::SeqCst) == KEY_COUNT);
        // 由于不同key之间的线程应当不互相影响, 所以总时间应该稍稍高于睡眠的时间.
        assert!(elapsed > SLEEP_DURATION && elapsed < SLEEP_DURATION * 2);

        // 所有执行完成后, 控制组内部应该是空的.
        assert!(control_group.is_empty());
    }

}
